import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt

"""
线性回归模型
"""

# 设置随机种子以确保结果可重现
torch.manual_seed(42)
# 1. 创建一些示例数据
# 生成带有一些噪声的线性数据
x = torch.linspace(0, 10, 100).reshape(-1, 1)
y = 2 * x + 1 + torch.randn(x.size()) * 2 # y = 2x + 1 + 噪声
print(f"数据形状: x={x.shape}, y={y.shape}")

# 2. 定义模型
class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(1, 1)   # 输入维度1，输出维度1

    def forward(self, x):
        return self.linear(x)

model = LinearRegression()
print(f"模型结构: {model}")
print(f"初始权重: {model.linear.weight.data}")
print(f"初始偏置: {model.linear.bias.data}")
# 3. 定义损失函数和优化器
criterion = nn.MSELoss() # 均方误差损失
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)   # 随机梯度下降
# 4. 训练模型
losses = []  # 记录损失值
epochs = 100

for epoch in range(epochs):
    # 前向传播
    prediction = model(x)
    loss = criterion(prediction, y)

    # 反向传播和优化
    optimizer.zero_grad()  # 清除之前的梯度
    loss.backward()        # 反向传播计算梯度
    optimizer.step()       # 更新

    # 记录损失
    losses.append(loss.item())

    # 每10个epoch打印一次损失
    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

# 5. 可视化结果
plt.figure(figsize=(12, 4))

# 绘制原始数据和拟合线
plt.subplot(1, 2, 1)
plt.scatter(x.numpy(), y.numpy(), label='原始数据', alpha=0.6)
plt.plot(x.numpy(), model(x).detach().numpy(), 'r-', label='拟合线')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title('线性回归拟合')

# 绘制损失曲线
plt.subplot(1, 2, 2)
plt.plot(losses)
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('训练损失')

plt.tight_layout()
plt.savefig('linear_regression_result.png')
plt.show()

# 6. 打印最终参数
print(f"\n训练后的权重: {model.linear.weight.data.item():.4f}")
print(f"训练后的偏置: {model.linear.bias.data.item():.4f}")
print(f"真实关系: y = 2x + 1")
print(f"学习到的关系: y = {model.linear.weight.data.item():.4f}x + {model.linear.bias.data.item():.4f}")